# Evaluate of SIFA
import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import numpy as np
import cv2

def create_visual_anno(anno):
    assert np.max(anno) < 7 # only 7 classes are supported, add new color in label2color_dict
    label2color_dict = {#BGR
        0: [0, 0, 0],
        1: [0, 0 , 255],  #红色
        2: [0, 255, 0], #绿色
        3: [122, 209, 255], #黄色
        4: [255, 215, 0],  #浅蓝色
        5: [160, 32, 100],
        6: [255, 64, 64],
        7: [139, 69, 19],
    }
    # visualize
    visual_anno = np.zeros((anno.shape[0], anno.shape[1], 3), dtype=np.uint8)
    for i in range(visual_anno.shape[0]):  # i for h
        for j in range(visual_anno.shape[1]):
            color = label2color_dict[anno[i, j]]
            visual_anno[i, j, 0] = color[0]
            visual_anno[i, j, 1] = color[1]
            visual_anno[i, j, 2] = color[2]

    return visual_anno

def norm_01(image):
    mn = np.min(image)
    mx = np.max(image)
    image = (image - mn) / (mx - mn).astype(np.float32)
    return image

def save_img(image):
    image = norm_01(image)
    image = (image * 255).astype(np.uint8)
    return image
def load_npz(path):
    img = np.load(path)['arr_0']
    gt = np.load(path)['arr_1']
    return img, gt

import os

import SimpleITK as sitk

def processAndSave(data_dir=r"D:\AllData\abdominalDATA\abdominalDATA\MR_T2_npy",
                   save_dir=r"D:\AllData\abdominalDATA\abdominalDATA\MRI3Dnii"):
    '''
    将sifa代码公开的npz数据集转化为3D的nii文件
    标注id和器官的对应关系：
    1：肝脏	2：右肾	3：左肾	4：脾
    '''
    if not os.path.exists(os.path.join(save_dir,"volume")):
        os.makedirs(os.path.join(save_dir,"volume"))
    if not os.path.exists(os.path.join(save_dir,"segmentation")):
        os.makedirs(os.path.join(save_dir,"segmentation"))
    number = 0

    subject_list=os.listdir(data_dir)
    subject_list.sort()
    for index,example in enumerate(subject_list):
        #单独对每一个3D数据进行z-score归一化
        image_array_list = []
        label_array_list = []
        npz_list=os.listdir(os.path.join(data_dir,example))
        for n in npz_list:
            name=os.path.join(data_dir,example,n)
            image_array,label_array=load_npz(name)
            image_array=save_img(image_array)
            image_array_list.append(image_array)
            label_array_list.append(label_array)
        #以整个3D图像进行z-score,并拼接image和label的每个2D切片，[image, mask] 。并保存
        image_array_list=np.array(image_array_list)
        image_array_list=(image_array_list-image_array_list.mean())/image_array_list.std()
        label_array_list=np.array(label_array_list)
        label_nii=sitk.GetImageFromArray(label_array_list)
        image_nii=sitk.GetImageFromArray(image_array_list)
        sitk.WriteImage(image_nii,os.path.join(save_dir,'volume',example+'.nii.gz'))
        sitk.WriteImage(label_nii, os.path.join(save_dir, 'segmentation', example + '.nii.gz'))
        print(index)
    print(number)
def getMaxAndMin(data_dir=r"D:\AllData\abdominalDATA\abdominalDATA\MR_T2_npy"):
    '''
    求整个数据集的最大值和最小值
    '''
    number = 0
    subject_list=os.listdir(data_dir)
    subject_list.sort()
    l=0#整个数据集的最大值
    m=0#整个数据集的最小值
    for index,example in enumerate(subject_list):
        #单独对每一个3D数据进行z-score归一化
        image_array_list = []
        label_array_list = []
        npz_list=os.listdir(os.path.join(data_dir,example))
        for n in npz_list:
            name=os.path.join(data_dir,example,n)
            image_array,label_array=load_npz(name)
            # image_array=save_img(image_array)
            image_array_list.append(image_array)
            label_array_list.append(label_array)

        #以整个3D图像进行z-score,并拼接image和label的每个2D切片，[image, mask] 。并保存
        image_array_list=np.array(image_array_list)
        #image_array_list=(image_array_list-image_array_list.mean())/image_array_list.std()
        label_array_list=np.array(label_array_list)

        l=max(l,image_array_list.max())
        m=min(m,image_array_list.min())

    print(l,m)
    return l,m
if __name__ == '__main__':
    # processAndSave(data_dir=r"D:\AllData\abdominalDATA\abdominalDATA\MR_T2_npy",
    #                save_dir=r"D:\AllData\abdominalDATA\abdominalDATA\MRI3Dnii")
    # processAndSave(data_dir=r"D:\AllData\abdominalDATA\abdominalDATA\CT_npy",
    #                save_dir=r"D:\AllData\abdominalDATA\abdominalDATA\CT3Dnii")
    getMaxAndMin(data_dir=r"/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/SIFAgf/BTCV_CT_npy")
    getMaxAndMin(data_dir=r"/home/liukai/AllData/AbdominalOrgansForDomain/ForSIFA/SIFAgf/CHAOS_MR_T2_npy")
    '''
    BTCV:  3.4539323 -1.2810926
    CHAOS: 4.1356735 -1.1623437
    

    '''
